"""
02_state_basics.py - Understanding State in LangGraph
=====================================================
Chapter 3: TypedDict, Pydantic, Annotations, and Reducers.

KEY CONCEPTS:
  ★ TypedDict - Simple state definition with type hints
  ★ Pydantic BaseModel - Production state with validation
  ★ Annotated - Adding reducers to control state updates
  ★ Reducer - Function that merges state updates (e.g., append vs replace)
  ★ MessagesState - Built-in state for chat applications

Run: python src/02_state_basics.py
"""

from typing import TypedDict, Annotated, List, Optional
from operator import add
from langgraph.graph import StateGraph, START, END


# ============================================================
# METHOD 1: TypedDict (Simple, great for learning)
# ============================================================

class SimpleState(TypedDict):
    """Basic state using TypedDict.
    
    Each field is a key in the state dictionary.
    Type hints tell LangGraph what type each field should be.
    No runtime validation - just type hints.
    """
    user_name: str
    greeting: str
    language: str


def generate_greeting(state: SimpleState) -> dict:
    """Generate a greeting. Returns ONLY the fields to update."""
    name = state["user_name"]
    lang = state.get("language", "English")
    
    greetings = {
        "English": f"Hello, {name}! Welcome to LangGraph!",
        "Spanish": f"Hola, {name}! Bienvenido a LangGraph!",
        "French": f"Bonjour, {name}! Bienvenue a LangGraph!",
        "Hindi": f"Namaste, {name}! LangGraph mein aapka swagat hai!",
    }
    
    return {"greeting": greetings.get(lang, greetings["English"])}


# ============================================================
# METHOD 2: State with Annotated Reducers
# ============================================================

class ChatState(TypedDict):
    """State with a reducer on messages.
    
    ★ KEY CONCEPT: Annotated[type, reducer]
    
    Without reducer: new value REPLACES old value
    With reducer:    reducer(old_value, new_value) determines result
    
    Annotated[list, add] means:
    - Field type is `list`
    - When updated, use Python's `add` (list concatenation)
    - {"messages": [new]} APPENDS to existing messages
    """
    messages: Annotated[list, add]  # ← APPEND reducer
    agent_name: str                  # ← No reducer = REPLACE


# ============================================================
# METHOD 3: Pydantic (Production-grade)
# ============================================================

try:
    from pydantic import BaseModel, Field
    from datetime import datetime
    
    class ProductionState(BaseModel):
        """Production state with Pydantic validation.
        
        ADVANTAGES over TypedDict:
        ✓ Runtime validation (catches bad data immediately)
        ✓ Default values
        ✓ Complex nested types
        ✓ Built-in serialization (JSON, dict)
        ✓ Clear error messages
        """
        messages: List[dict] = Field(
            default_factory=list,
            description="Conversation message history"
        )
        current_agent: str = Field(
            default="router",
            description="Which agent is currently active"
        )
        context: Optional[str] = Field(
            default=None,
            description="Retrieved context for RAG"
        )
        iteration_count: int = Field(
            default=0,
            ge=0,
            le=10,
            description="Number of iterations (max 10)"
        )
    
    PYDANTIC_AVAILABLE = True
except ImportError:
    PYDANTIC_AVAILABLE = False


# ============================================================
# DEMONSTRATIONS
# ============================================================

def demo_simple_state():
    """Demo 1: Simple TypedDict state."""
    print("\n📋 Demo 1: Simple TypedDict State")
    print("-" * 40)
    
    graph = StateGraph(SimpleState)
    graph.add_node("greeter", generate_greeting)
    graph.add_edge(START, "greeter")
    graph.add_edge("greeter", END)
    app = graph.compile()
    
    for lang in ["English", "Spanish", "Hindi"]:
        result = app.invoke({"user_name": "Developer", "language": lang})
        print(f"  {lang}: {result['greeting']}")


def demo_reducer_state():
    """Demo 2: Annotated reducer state."""
    print("\n📋 Demo 2: Annotated Reducer State")
    print("-" * 40)
    
    def add_message(state: ChatState) -> dict:
        """Add a message. Thanks to the reducer, this APPENDS."""
        return {
            "messages": [f"Message #{len(state['messages']) + 1}"],
            "agent_name": "chatbot"
        }
    
    graph = StateGraph(ChatState)
    graph.add_node("add_msg", add_message)
    graph.add_edge(START, "add_msg")
    graph.add_edge("add_msg", END)
    app = graph.compile()
    
    # First invocation
    result = app.invoke({"messages": ["Hello!"], "agent_name": ""})
    print(f"  After invoke 1: {result['messages']}")
    # Output: ["Hello!", "Message #2"] - APPENDED!


if __name__ == "__main__":
    print("=" * 60)
    print("02: Understanding State in LangGraph")
    print("=" * 60)
    
    demo_simple_state()
    demo_reducer_state()
    
    if PYDANTIC_AVAILABLE:
        print("\n📋 Demo 3: Pydantic State")
        print("-" * 40)
        state = ProductionState()
        print(f"  Default state: {state.model_dump()}")
        
        state.iteration_count = 5
        print(f"  Updated: iteration_count = {state.iteration_count}")
        
        try:
            state.iteration_count = 15  # Exceeds max!
            print("  Should not reach here")
        except Exception as e:
            print(f"  Validation caught: {type(e).__name__}")
    
    print(f"\n{'=' * 60}")
    print("✅ State fundamentals complete!")
    print("=" * 60)
